"""
验证入口 - 验证已训练的模型
"""
import argparse
import torch
from pathlib import Path

from configs.default import DefaultConfig
from models import build_model
from core.validator import Validator
from utils.checkpoint import load_checkpoint
from utils.general import set_seed, colorstr, select_device


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Validate Meniscus Volume Prediction Model')
    
    parser.add_argument('--weights', type=str, required=True,
                       help='path to model weights')
    parser.add_argument('--data', type=str, default='/root/workspace',
                       help='data directory')
    parser.add_argument('--batch-size', type=int, default=32,
                       help='batch size')
    parser.add_argument('--workers', type=int, default=8,
                       help='number of workers')
    parser.add_argument('--device', type=str, default='',
                       help='cuda device')
    parser.add_argument('--compare', action='store_true',
                       help='show prediction vs target comparison')
    parser.add_argument('--num-samples', type=int, default=10,
                       help='number of samples to display in comparison')
    
    return parser.parse_args()


def main():
    """主函数"""
    # 解析参数
    args = parse_args()
    
    # 检查权重文件
    weights_path = Path(args.weights)
    if not weights_path.exists():
        print(colorstr('red', f'Error: Weights file not found: {weights_path}'))
        return
    
    print(colorstr('bright_blue', 'bold', '\n' + '='*70))
    print(colorstr('bright_blue', 'bold', 'Model Validation'.center(70)))
    print(colorstr('bright_blue', 'bold', '='*70))
    
    # 加载checkpoint获取配置
    checkpoint = torch.load(weights_path, map_location='cpu')
    
    # 创建配置
    config = DefaultConfig()
    if 'config' in checkpoint:
        for k, v in checkpoint['config'].items():
            if hasattr(config, k):
                setattr(config, k, v)
    
    # 更新配置
    config.data_dir = args.data
    config.batch_size = args.batch_size
    config.num_workers = args.workers
    config.device = select_device(args.device)
    
    # 打印信息
    print(f"\nWeights: {weights_path}")
    print(f"Model Type: {config.model_type}")
    print(f"Device: {config.device}")
    
    # 构建模型
    print(colorstr('bright_green', '\nBuilding model...'))
    model = build_model(config).to(config.device)
    
    # 创建验证器
    validator = Validator(model, config)
    
    # 加载权重
    validator.load_model(weights_path)
    
    # 验证
    print(colorstr('bright_green', '\nValidating...'))
    val_loss, metrics, predictions, targets = validator.validate()
    
    # 显示对比（可选）
    if args.compare:
        validator.compare_predictions(predictions, targets, args.num_samples)
    
    print(colorstr('bright_green', 'bold', '\n✓ Validation completed!'))


if __name__ == '__main__':
    main()